function [x_opt,a_opt] = optimization_torso
global ns_time s_time ns_slope s_knee ns_knee hip_pos pos_com


%%%%%%%%%%%%%Add paths to mathematica compiled files:

addpath('./build_torso')
addpath('./wrappers_torso')


addpath('./buildopt_torso')
addpath('./Data')


%%%%%%%%%%%%%Loading the data:

load('fred_data_hippos.mat')
load('fred_data_nsslope.mat')
load('fred_data_sknee.mat')
load('fred_data_nsknee.mat')
load('fred_data_torso.mat')

% Hack here to replace COM pos (ryan sinnet)
pos_com = torso;

%%%%%%%%%%%%%Outputs from loading the data:
HipData = [ns_time, hip_pos];
NSslopeData = [ns_time,ns_slope];
SkneeData = [s_time, s_knee];
NSkneeData = [ns_time,ns_knee];
PcomData = [ns_time, pos_com];

%%%%%%%%%%%%%Initial fit parameters (as determined by mathemaitca right now):
%%%%%%%%%%%%%% TODO:  determine the initial parameters from the data in
%%%%%%%%%%%%%% matlab

% a0 = [    1.1969         0         0         0         0
%     0.2021    7.8404    0.2248   -0.8070    0.1871
%    -0.0828   13.3160    0.2073    4.1548    0.2574
%    -0.3809  -10.9795   -0.1967   -0.4214    0.6585
%    -3.2612   -0.0000    5.8197    0.3775    3.0234];

load('Afit.mat', 'a0');
   
  
    
%%%%%%%%%%%%%% Minimization to verify optimality of initial parameters w.r.t the HBC
% options = optimset ('Display','iter','MaxFunEvals',6000);
% [a_fit,cost_fit] = fminunc(@cost,flatten(a0),options)


%%%%%%%%%%%%%% In the case that you don't want to run the optimizatino to
%%%%%%%%%%%%%% verify the parameters, uncomment this:
 a_fit = flatten(a0);

%%%%%%%%%%%%%% Constrain the parameters of the hip position to be zero
ubmat = Inf*ones(5,5);
lbmat = -Inf*ones(5,5);
ubmat(1,2) = 0; ubmat(1,3) = 0; ubmat(1,4) = 0; ubmat(1,5) = 0; 
lbmat(1,2) = 0; lbmat(1,3) = 0; lbmat(1,4) = 0; lbmat(1,5) = 0;
ub = flatten(ubmat);
lb = flatten(lbmat);

%%%%%%%%%%%%%% Other conditions for the optimization:
Aeq = [];
beq = [];
A = [];
b = [];

%%%%%%%%%%%%%% Optimizing to achieve HZD with the HBC:

%%%%% Conditions for optimization: tight error 
 errorbound = 1e-5;
options = optimset('Algorithm','sqp',...
    'Display','iter','TolX',errorbound,...
'MaxFunEvals',4000,'TolCon',errorbound);

%%%%% Conditions for optimization:
% options = optimset ('Display','iter','MaxFunEvals',6000);

%%%%% Optimization
tic
[a_out,cost_out] = fmincon(@cost,a_fit,A,b,Aeq,beq,lb,ub,@mycon,options);
toc

%%%%%%%%%%%%%% Output the parameter matrix and the fit matrix:
a_opt = unflatten(a_out)
a_fit = unflatten(a_fit)

%%%%%%%%%%%%%% Output the initial condition:
q = inverse_kin(a_opt);
qdot = H_minus(q,a_opt)\[a_opt(1,1); 0; 0; 0; 0];
x_opt = [q; qdot]

%%%%%%%%%%%%%% Output the final time: 
disp(['Step time: ' tau(x_opt,a_opt)]);


%%%%%%%%%%%%%% Save initial and final condition
save('x_opt.mat','x_opt')
save('a_opt.mat','a_opt')

%%%%%%%%%%%%%% Plot the output vs. fit vs. the human data

figure(10); clf; 
subplot(2,3,1); plot(ns_time,yd_hip_pos(ns_time,a_opt),'b',ns_time,yd_hip_pos(ns_time,a_fit),'r',ns_time,hip_pos,'k.') 
title('Hip Position')
xlabel('time (s)');
% ylabel('Desired vs. Fit vs. Data')
legend({'Desired','Fit','Data'});
subplot(2,3,2); plot(ns_time,yd_ns_slope(ns_time,a_opt),'b',ns_time,yd_ns_slope(ns_time,a_fit),'r',ns_time,ns_slope,'k.')
title('Non-Stance Slope')
xlabel('time (s)');
% ylabel('Desired vs. Fit vs. Data')
legend({'Desired','Fit','Data'});
subplot(2,3,4); plot(s_time,yd_s_knee(s_time,a_opt),'b',s_time,yd_s_knee(s_time,a_fit),'r',s_time,s_knee,'k.')
title('Stance Knee')
xlabel('time (s)');
% ylabel('Desired vs. Fit vs. Data')
legend({'Desired','Fit','Data'});
subplot(2,3,5); plot(ns_time,yd_ns_knee(ns_time,a_opt),'b',ns_time,yd_ns_knee(ns_time,a_fit),'r',ns_time,ns_knee,'k.')
title('Non-Stance Knee')
xlabel('time (s)');
% ylabel('Desired vs. Fit vs. Data')
legend({'Desired','Fit','Data'});
subplot(2,3,3); plot(ns_time,yd_pos_com(ns_time,a_opt),'b',ns_time,yd_pos_com(ns_time,a_fit),'r',ns_time,pos_com,'k.')
title('Position of the COM')
xlabel('time (s)');
% ylabel('Desired vs. Fit vs. Data')
legend({'Desired','Fit','Data'});


%%%%%%%%%%%%%% Run main or step
  main_AMBER_torso
% step(a_opt)

end



function ret = cost(aflat)
global ns_time s_time ns_slope s_knee ns_knee hip_pos pos_com

%%%%%%%%%%%%%% Unflatten the a matrix
a = unflatten(aflat);

%%%%%%%%%%%%%% Cost for each canonical function
% cost_ns_slope = norm(ns_slope-yd_ns_slope(ns_time,a))^2;
% cost_hip_pos = norm(hip_pos-yd_hip_pos(ns_time,a))^2;
% cost_s_knee = norm(s_knee-yd_s_knee(s_time,a))^2;
% cost_ns_knee = norm(ns_knee-yd_ns_knee(ns_time,a))^2;
% cost_pos_com = norm(pos_com-yd_pos_com(ns_time,a))^2;

%%%%%%%%%%%%%% Cost for each canonical function

cost_ns_slope = sum(abs(ns_slope-yd_ns_slope(ns_time,a)));
cost_hip_pos = sum(abs(hip_pos-yd_hip_pos(ns_time,a)));
cost_s_knee = sum(abs(s_knee-yd_s_knee(s_time,a)));
cost_ns_knee = sum(abs(ns_knee-yd_ns_knee(ns_time,a)));
cost_pos_com = sum(abs(pos_com-yd_pos_com(ns_time,a)));

%%%%%%%%%%%%%% Wieghting for the fits

beta_ns_slope = 1/(max(ns_slope)-min(ns_slope));
beta_hip_pos = 1/(max(hip_pos)-min(hip_pos));
beta_s_knee = 1/(max(s_knee)-min(s_knee));
beta_ns_knee = 1/(max(ns_knee)-min(ns_knee));
beta_pos_com = 1/(max(pos_com)-min(pos_com));
beta = beta_ns_slope+beta_hip_pos+beta_s_knee+beta_ns_knee+beta_pos_com;

cost_all = [
      beta_ns_slope*cost_ns_slope,...
      beta_hip_pos*cost_hip_pos,...
      beta_s_knee*cost_s_knee,...
      beta_ns_knee*cost_ns_knee,...
      beta_pos_com*cost_pos_com
      ];

% cost_all = [
%       cost_ns_slope,...
%       cost_hip_pos,...
%       cost_s_knee,...
%       cost_ns_knee,...
%       cost_pos_com...
%       ];
  
%%%%%%%%%%%%%% Return the cost

%  ret = max(cost_all);
% ret = max(cost_all)/beta;
%    ret = sum(cost_all);
  ret = sum(cost_all)/beta;

end


%%%%%%%%%%%%%% Constraints to enforce HZD (or partial HZD)

function [c,ceq] = mycon(aflat)
global R1

%%%%%%%%%%%%%% Unflatten the a matrix
a = unflatten(aflat);

%%%%%%%%%%%%%% Determine the initial condition (on the guard) for the
%%%%%%%%%%%%%% current parameters
q = inverse_kin(a);
% q = find_inverse_kin(a);
qdot = H_minus(q,a)\[a(1,1); 0; 0; 0; 0];
x = [q; qdot];

ic = resetFunc(x);
p_hip_ic = phip_sca(ic);

% clear h t
% Tfinal = max(.2,tau(q,a));
% t = 0:.01:abs(Tfinal);
% h = zeros(length(t), 1);
% for i = 1:length(t)
%     h(i) = height(t(i),a,p_hip_ic);
% end
% 
% [hmin,hminindex] = min(h(3:end-2));
% Tmin = t(hminindex);
% 
% figure(20)
% plot(t,h,'b',...
%     t,zeros(1,length(t)),'k',...
%     Tfinal,height(Tfinal,a,p_hip_ic),'gs',...
%     Tmin,hmin,'ro')
% axis([0 .6 -0.02 .2])


% 
% Tmin = fminbnd(@(t)height(t,a,p_hip_ic),0,tau(q,a));
% hmin = height(Tmin,a,p_hip_ic);

Ttau = tau(q,a);
if imag(Ttau)>0||Ttau < 0
Ttau = 1;
end

%%%%%%%%%%%%%% H is the post impact velocities. 
H = H_plus(q,a)*P(x);

    %%%%%%%%%%%%%% Require the post impact velocities be completlty determined:
      % Hzd = H - [a(1,1); 0; 0; 0; 0];
     Hzd = H(2:5,:) - [0; 0; 0; 0];

    %%%%%%%%%%%%%% Pre and Post Impact knee angles 
    
R1 = [1  1  1 -1 -1;
      0  0  0  0  1;
      0  0  0  1  0;
      0  0  1  0  0;
      0  1  0  0  0];
  
     qR1 = R1*q;   
     
    %%%%%%%%%%%%%% Nonlinear inequalities
    %%%%%%%%%%%%%%%%% gaurd must be negative
    %%%%%%%%%%%%%%%%% Time must be less than the human time
    c = [
           h_dot_minus(q)*qdot+.01;
 %         tau(q,a)-0.5208;
%             -hmin + .004;
%           Hfullzd-bound;
%          -Hfullzd+bound;
%         -qR1(2,1)+.000;
%         -q(2,1)+.000;
%         -qR1(5,1)+.000;
%         -q(5,1)+.000;
%          q(4,1)-.2;
%         qR1(4,1)-.2;
%         qR1(2,1)-.2;
%         q(2,1)-.2
        ];
    
%%%%%%%%%%%%%% Nonlinear equalities
%%%%%%%%%%%%%%%%% y pre impact must be zero
%%%%%%%%%%%%%%%%% y dot post impact must be zero
%%%%%%%%%%%%%%%%% guard must be satisfied
ceq = [
       y_minus(q,a);
       Ttau-0.5208;
       Hzd];

   
end


%%%%%%%%%%%%%% Computing the change in impact velocities

function ret = P(x)
      
dimq = length(x)/2;
xplus = resetFunc(x);

ret = xplus(dimq+1:end);    
        
end     

function ret = inverse_kin(a)
addpath('./buildopt_torso')




%%%%If using the old version (no torso) use this:

% theta_a(a)
% 

%%%%If using the new version with the torso, use this: 
%%%%%(the calculation is broken up, saves a lot of computation time)

q = zeros(5,1);
q(2) = theta_a2(q,a);
q(5) = theta_a5(q,a);
q(1) = theta_a1(q,a);
q(4) = theta_a4(q,a);
q(3) = theta_a3(q,a);

ret = q;




end


%%%%%%%%%%%%%% Canonical human functions

function ret = yd_ns_slope(t,a)

ret =exp(-a(2,4)*t).*(a(2,1)*cos(a(2,2)*t)+a(2,3)*sin(a(2,2)*t))+a(2,5);

end

function ret = yd_hip_pos(t,a)

ret = a(1,1)*t;

end

function ret = yd_s_knee(t,a)

ret = exp(-a(3,4)*t).*(a(3,1)*cos(a(3,2)*t)+a(3,3)*sin(a(3,2)*t))+a(3,5);

end

function ret = yd_ns_knee(t,a)

ret = exp(-a(4,4)*t).*(a(4,1)*cos(a(4,2)*t)+a(4,3)*sin(a(4,2)*t))+a(4,5);

end

function ret = yd_pos_com(t,a)

ret = exp(-a(5,4)*t).*(a(5,1)*cos(a(5,2)*t)+a(5,3)*sin(a(5,2)*t))+a(5,5);

end

function aflat = flatten(a)
aflat = a(:);

end

function aunflat = unflatten(aflat)
aunflat = reshape(aflat,5,5);

end


% 
% 
% function ret = u_calc(q, p, a, ep)
% 
% 
% ndof = length(q)/2;
% dx = q(ndof+1:end);
% % q
% % pause(1000)
% M    = D_mat(q);
% C    = C_mat(q);
% G    = G_vec(q);
% 
% % pause(0.5);
% 
% %%%Outputs
% y_a1     = ya1_sca(q, p, a);
% y_a2     = ya2_vec(q, p, a);
% y_d1     = yd1_sca(q, p,  a);
% y_d2     = yd2_vec(q, p,  a);
% 
% %%%Outputs
% y1 = y_a1 - y_d1;
% y2 = y_a2 - y_d2;
% 
% 
% y_a = [y_a1; y_a2];
% y_d = [y_d1; y_d2];
% 
% %%%Jacobians of Outputs
% Dy_a1    = Dya1_mat(q, p, a);
% Dy_d1    = Dyd1_mat(q, p, a);
% Dy_a2    = Dya2_mat(q, p, a);
% Dy_d2    = Dyd2_mat(q, p, a);
% 
% Dy_1 = Dy_a1 - Dy_d1;
% Dy_2 = Dy_a2 - Dy_d2;
% 
% %%%Control Fields
% vf    = [dx; M \ (-C*dx - G)];
% B_IO  = eye(ndof);
% gf    = [zeros(size(B_IO)); M \ B_IO];
% 
% %%%Lie Derivatives
% Lgy1 = Dy_1*gf;
% Lgy2 = Dy_2*gf;
% Lfy1 = Dy_1*vf;
% Lfy2 = Dy_2*vf;
% 
% %%%Second Order Jacobians
% 
% DLfy_a2  = DLfya2_mat(q, p, a);
% DLfy_d2  = DLfyd2_mat(q, p, a);
% 
% DLfy2 = DLfy_a2-DLfy_d2;
% 
% %%%Second Lie Derivatives
% 
% LfLfy2 = DLfy2*vf;
% LgLfy2 = DLfy2*gf;
% 
% % tp = [tp, t+tp0];
% % yact = [yact, y_a];
% % ydes = [ydes, y_d];
% 
% % Vector field
% 
% A = [Lgy1; LgLfy2];
% 
% u = -A \ ([0; LfLfy2]+[Lfy1; 2*ep*Lfy2]+ [ep*y1; ep^2*y2]);
% 
% 
% % U = [U, u];
% % Pow = [Pow, u.*dx];
% 
% ret = u;
% 
% end

